#include "QBasicOp.h"
#include "BasicOp.h"
#include <assert.h>

using namespace DL;

Tensor3D<int8>& DL::QConv2D(const Tensor3D<int8>& tensor, QATParams in, int stride, int padding, const Tensor4D<int8>& weight, const Tensor1D<float>& scale, const Tensor1D<float>& bias, QATParams out)
{
	int outWidth = (tensor.width - weight.width + 2 * padding) / stride + 1;
	int outHeight = (tensor.height - weight.height + 2 * padding) / stride + 1;
	Tensor3D<int8>* outTensor = new Tensor3D<int8>(outWidth, outHeight, weight.deepth);

	for (int b = 0; b < outTensor->channel; b++)
	{
		float perChannelScale = in.scale * scale.Get(b);
		for (int y = 0; y < outTensor->height; y++)
		{
			for (int x = 0; x < outTensor->width; x++)
			{
				int result = 0;
				for (int c = 0; c < tensor.channel; c++)
				{
					for (int ky = 0; ky < weight.height; ky++)
					{
						for (int kx = 0; kx < weight.width; kx++)
						{
							result += tensor.Get(x * stride - padding + kx, y * stride - padding + ky, c) * weight.Get(kx, ky, c, b);
						}
					}
				}
				float resultQta = result * perChannelScale;
				resultQta += bias.Get(b);
				resultQta /= out.scale;
				outTensor->Set(x, y, b, (int8)resultQta);
			}
		}
	}
	return *outTensor;
}

Tensor3D<int8>& DL::QMaxPooling2D(const Tensor3D<int8>& tensor, int size, int stride, int padding)
{
	int outWidth = (tensor.width - size + 2 * padding) / stride + 1;
	int outHeight = (tensor.height - size + 2 * padding) / stride + 1;
	Tensor3D<int8>* outTensor = new Tensor3D<int8>(outWidth, outHeight, tensor.channel);

	for (int y = 0; y < outTensor->height; y++)
	{
		for (int x = 0; x < outTensor->width; x++)
		{
			for (int c = 0; c < tensor.channel; c++)
			{
				int8 result = SCHAR_MIN;
				for (int ky = 0; ky < size; ky++)
				{
					for (int kx = 0; kx < size; kx++)
					{
						int8 value = tensor.Get(x * stride - padding + kx, y * stride - padding + ky, c);
						result = result > value ? result : value;
					}
				}
				outTensor->Set(x, y, c, result);
			}
		}
	}
	return *outTensor;
}

Tensor3D<int8>& DL::QRelu(const Tensor3D<int8>& tensor)
{
	Tensor3D<int8>* outTensor = new Tensor3D<int8>(tensor.width, tensor.height, tensor.channel);

	for (int y = 0; y < outTensor->height; y++)
	{
		for (int x = 0; x < outTensor->width; x++)
		{
			for (int c = 0; c < tensor.channel; c++)
			{
				int8 value = tensor.Get(x, y, c);
				value = value > 0 ? value : 0;
				outTensor->Set(x, y, c, value);
			}
		}
	}
	return *outTensor;
}

Tensor1D<int8>& DL::QRelu(const Tensor1D<int8>& tensor)
{
	Tensor1D<int8>* outTensor = new Tensor1D<int8>(tensor.width);
	for (int x = 0; x < outTensor->width; x++)
	{
		int8 value = tensor.Get(x);
		value = value > 0 ? value : 0;
		outTensor->Set(x, value);
	}
	return *outTensor;
}

Tensor1D<int8>& DL::QFlatten(const Tensor3D<int8>& tensor)
{
	Tensor1D<int8>* outTensor = new Tensor1D<int8>(tensor.width * tensor.height * tensor.channel);
	for (int c = 0; c < tensor.channel; c++)
	{
		for (int y = 0; y < tensor.height; y++)
		{
			for (int x = 0; x < tensor.width; x++)
			{
				int8 value = tensor.Get(x, y, c);
				outTensor->Set(c * tensor.framePiexlCount + y * tensor.width + x, value);
			}
		}
	}
	return *outTensor;
}

Tensor1D<int8>& DL::QLinear(const Tensor1D<int8>& tensor, QATParams in, const Tensor2D<int8>& weight, const Tensor1D<float>& scale, const Tensor1D<float>& bias, QATParams out)
{
	Tensor1D<int8>* outTensor = new Tensor1D<int8>(weight.height);
	for (int outC = 0; outC < outTensor->width; outC++)
	{
		int result = 0;
		for (int inC = 0; inC < tensor.width; inC++)
		{
			result += tensor.Get(inC) * weight.Get(inC, outC);
		}
		float resultQta = result * in.scale * scale.Get(outC);
		resultQta += bias.Get(outC);
		resultQta /= out.scale;
		outTensor->Set(outC, (int8)resultQta);
	}
	return *outTensor;
}

Tensor1D<float>& DL::Dequantify(const Tensor1D<int8>& tensor, QATParams qat)
{
	Tensor1D<float>* outTensor = new Tensor1D<float>(tensor.width);
	for (int i = 0; i < tensor.width; i++)
	{
		outTensor->Set(i, tensor.Get(i) * qat.scale);
	}
	return *outTensor;
}

Tensor3D<int8>& DL::Quantify(const Tensor3D<float>& tensor, QATParams qat)
{
	Tensor3D<int8>* outTensor = new Tensor3D<int8>(tensor.width, tensor.height, tensor.channel);
	for (int c = 0; c < tensor.channel; c++)
	{
		for (int y = 0; y < tensor.height; y++)
		{
			for (int x = 0; x < tensor.width; x++)
			{
				int8 value = tensor.Get(x, y, c) / qat.scale;
				outTensor->Set(x, y, c, value);
			}
		}
	}
	return *outTensor;
}



